from torch.utils.data import Dataset
from torch.utils.data.dataset import T_co
from utils.plus_function import get_bin_str_list_range_n
from utils.raw_data1 import RawData
from exp05.instance import SentenceInstance, TrainingInstance, TestInstance


class TrainDataSet(Dataset):
    def __init__(self, raw_data: RawData):
        self.meta_instances = raw_data
        self.training_instances = []

        def generate_train_instance_list(meta_ins: SentenceInstance):
            n_triplet = len(meta_ins.triplets)
            bin_str_list = get_bin_str_list_range_n(n_triplet)
            return [TrainingInstance(meta_ins, bin_str) for bin_str in bin_str_list]
            pass

        for meta_instance in self.meta_instances:
            self.training_instances.extend(generate_train_instance_list(meta_instance))
            pass
        pass

    def __getitem__(self, index) -> T_co:
        return self.training_instances[index]

    def __iter__(self):
        for elem in self.training_instances:
            yield elem
            pass
        pass

    def __len__(self) -> int:
        return len(self.training_instances)

    pass


class TestDataSet(Dataset):
    def __init__(self, raw_data: RawData):
        self.meta_instances = raw_data
        self.test_instances = [TestInstance(v) for v in self.meta_instances]
        pass

    def __getitem__(self, index) -> T_co:
        return self.test_instances[index]

    def __iter__(self):
        for elem in self.test_instances:
            yield elem
            pass
        pass

    def __len__(self):
        return len(self.test_instances)

    pass
